import re
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, TypeVar
from warnings import warn

import rich


@dataclass
class Span:
    start: int
    end: int


@dataclass(frozen=True, kw_only=True)
class SourceSpan:
    span: Span


@dataclass(frozen=True, kw_only=True)
class SourceRaw:
    raw: bytes

    def __bytes__(self) -> bytes:
        return self.raw


@dataclass(frozen=True, kw_only=True)
class SourceChunk(SourceSpan, SourceRaw):
    pass


@dataclass(frozen=True, kw_only=True)
class Label:
    id: bytes

    def __bytes__(self) -> bytes:
        return b'/* @@ %b */' % self.id


@dataclass(frozen=True, kw_only=True)
class SourceLabel(SourceChunk, Label):
    pass


@dataclass(frozen=True, kw_only=True)
class LabeledAlert:
    id: bytes
    kind: bytes
    error: bytes

    def __bytes__(self) -> bytes:
        return b'/* @@@ %b %b %b */' % (
            self.id,
            self.kind,
            self.error,
        )


@dataclass(frozen=True, kw_only=True)
class SourceLabeledAlert(SourceChunk, LabeledAlert):
    pass


@dataclass(frozen=True, kw_only=True)
class AlertPosition:
    file: bytes | None
    line: int
    column: int

    def __bytes__(self):
        result = [
            str(self.line).encode(),
            str(self.column).encode(),
        ]
        if self.file:
            result.insert(0, self.file)
        return b':'.join(result)


@dataclass(frozen=True, kw_only=True)
class PositionAlert:
    pos: AlertPosition
    kind: bytes
    error: bytes

    def __bytes__(self):
        return b'/* @@@ %b %b %b */' % (
            bytes(self.pos),
            self.kind,
            self.error,
        )


@dataclass(frozen=True, kw_only=True)
class SourcePositionAlert(SourceChunk, PositionAlert):
    pass


_ID = rb'\s+(?P<id>\w+)'
_ID_ERROR = rb'@\s*(?P<id_ref>\w+)\s+(?P<id_type>\S+)\s+(?P<id_error>.*?)'
_POS_ERROR = rb'[?]\s*(?:(?P<file>\S+):)?(?P<line>\d+):(?P<column>\d+)\s*(?P<pos_type>\S+)\s+(?P<pos_error>.*?)'
_SKIP = rb'#\s*(?:(?P<option>\w+)\s*=\s*(?P<value>\w+))'
_OTHER = rb'(?P<other>.*?)'

_INSIDE = b'|'.join(
    [
        b'(?:%b)' % _i
        for _i in [
            _ID,
            _ID_ERROR,
            _POS_ERROR,
            _SKIP,
            _OTHER,
        ]
    ]
)

_PATTERN = re.compile(
    pattern=rb'(?P<all>/\*\s*@@(?P<INSIDE>%b)\s*\*/)' % (_INSIDE),
    flags=re.DOTALL,
)


def _parse_text(test_source: bytes) -> Iterator[SourceChunk]:
    for m in _PATTERN.finditer(test_source):
        raw = bytes(m['all'])
        if other := m.group('other'):
            warn(f'Not parsed: {m.span()}\n{other!r}')
        elif id := m.group('id'):
            yield SourceLabel(
                raw=raw,
                span=Span(*m.span()),
                id=id,
            )
        elif m.group('id_ref'):
            yield SourceLabeledAlert(
                raw=raw,
                span=Span(*m.span()),
                id=m['id_ref'],
                kind=m['id_type'],
                error=m['id_error'],
            )
        elif m.group('line'):
            file = m.group('file')
            yield SourcePositionAlert(
                raw=raw,
                span=Span(*m.span()),
                pos=AlertPosition(
                    file=file if file else None,
                    line=int(m['line']),
                    column=int(m['column']),
                ),
                kind=m['pos_type'],
                error=m['pos_error'],
            )
        else:
            rich.print(m.span(), m.group())
            warn(f'Not parsed: {m.span()}\n{m}')


# return (line, column) with indexing from 1
def calc_position(data: bytes, pos: int) -> tuple[int, int]:
    return (
        data.count(b'\n', 0, pos) + 1,  # line
        pos - data.rfind(b'\n', 0, pos),  # column
    )


TSourceCunck = TypeVar('TSourceCunck', bound=SourceChunk)


@dataclass(frozen=True)
class Source:
    file: Path
    raw: bytes
    chunks: list[SourceChunk]

    def calc_position(self, pos: int) -> AlertPosition:
        line, column = calc_position(self.raw, pos)
        return AlertPosition(
            file=self.file.name.encode(),
            line=line,
            column=column,
        )

    def filter(self, t: type[TSourceCunck]) -> Iterator[TSourceCunck]:
        for chunk in self.chunks:
            if isinstance(chunk, t):
                yield chunk

    def find_labels(self, id: bytes) -> Iterator[SourceLabel]:
        for lbl in self.filter(SourceLabel):
            if lbl.id == id:
                yield lbl

    def label_position(self, id: bytes):
        for lbl in self.find_labels(id):
            return lbl.span.end
        return None

    def write(self, buffer: BinaryIO):
        for i in self.chunks:
            buffer.write(bytes(i))
        buffer.flush()

    def raw_write(self, buffer: BinaryIO):
        for i in self.chunks:
            buffer.write(i.raw)
        buffer.flush()

    def peek(self, pos: int):
        for m in self.chunks:
            if m.span.start <= pos and m.span.end > pos:
                return m
        return None


def _parse(source: bytes) -> Iterator[SourceChunk]:
    pos: int = 0
    for meta in _parse_text(source):
        start, end = (meta.span.start, meta.span.end)
        if start > pos:
            yield SourceChunk(
                span=Span(pos, start),
                raw=source[pos:start],
            )
        yield meta
        pos = end


def parse_source(source_file: Path):
    raw = source_file.read_bytes()
    alerts = [*_parse(raw)]
    return Source(
        file=source_file,
        raw=raw,
        chunks=alerts,
    )
